import time
import traceback
from copy import deepcopy
from typing import List, Dict

import openai

from conf.config import config, logger


session_list: Dict[str, List[Dict]] = dict()


class OpenAIBot(object):
    def __init__(self):
        if config.get("open_ai_proxy_server_enable"):  # 使用第三方代理服务器
            openai.api_key = config.get('open_ai_proxy_api_key')
            openai.api_base = config.get("open_ai_proxy_api_base")
        elif config.get("proxy_enable"):  # 使用本地代理
            openai.api_key = config.get('open_ai_api_key')
            openai.proxy = {"http": f"{config.get('proxy_host')}:{config.get('proxy_port')}", "https": f"{config.get('proxy_host')}:{config.get('proxy_port')}"}  # api被墙时，开启代理
        else:  # 无需翻墙
            openai.api_key = config.get('open_ai_api_key')

    def reply(self, question, session_id=None):

        question = question.strip()
        logger.info("[OPEN_AI QUERY] session_id={}, question={}".format(session_id, question))

        messages = Session.build_session_query(question, session_id)
        logger.info(f"messages = {messages}")
        logger.info(f"生成响应中...")
        answer = self.reply_text(messages=messages, retry_count=0)
        logger.info("[OPEN_AI REPLY] session_id={}, answer={}".format(session_id, answer))

        if session_id and question and answer:  # 有会话id时，才需要存储会话，否则认为是单轮对话
            Session.save_session(question, answer, session_id)
        return answer

    def reply_text(self, messages: List, retry_count=0):
        try:
            response = openai.ChatCompletion.create(
                model=config.get("open_ai_model"),  # 对话模型的名称
                messages=messages
            )
            response_content = response['choices'][0]['message']['content']
            return response_content
        except openai.error.RateLimitError:
            # rate limit exception
            error_str = traceback.format_exc()
            logger.error(error_str)

            if retry_count < 2:
                time.sleep(5)
                logger.warn("[OPEN_AI] RateLimit exceed, 第{}次重试".format(retry_count+1))
                return self.reply_text(messages, retry_count+1)
            else:
                return "提问太快啦，请休息一下再问我吧"
        except openai.error.InvalidRequestError:
            error_str = traceback.format_exc()
            logger.error(error_str)
            return "非法请求，请检查"
        except Exception:
            # unknown exception
            error_str = traceback.format_exc()
            logger.error(error_str)
            return "请再问我一次吧"

    def reply_img(self, query, retry_count=0):
        pass


class Session(object):
    global session_list

    @staticmethod
    def build_session_query(question, session_id):
        """
        结合历史会话，生成查询内容
        :param question: 查询的问题
        :param session_id: 会话id，区分不同会话
        :return: 含历史会话记录的查询内容
        """

        prompt: List = config.get("open_ai_prompt", [])
        session: List = session_list.get(session_id, [])

        session_query: List = deepcopy(prompt)

        # 添加历史会话记录
        if session:
            for conversation in session:
                session_query.append({"role": "user", "content": f"{conversation['question']}"})
                session_query.append({"role": "assistant", "content": f"{conversation['answer']}"})

        session_query.append({"role": "user", "content": f"{question}"})

        return session_query

    @staticmethod
    def save_session(question, answer, user_id):

        max_tokens = config.get("session_max_tokens")
        if not max_tokens:
            max_tokens = 2048
        conversation = dict()
        conversation["question"] = question
        conversation["answer"] = answer
        session = session_list.get(user_id)

        if session:
            session.append(conversation)
        else:
            # create session
            session_list[user_id] = [conversation]

        # discard exceed limit conversation
        Session.discard_exceed_conversation(session_list[user_id], max_tokens)

    @staticmethod
    def discard_exceed_conversation(session, max_tokens):
        count = 0
        count_list = list()
        for i in range(len(session)-1, -1, -1):
            # count tokens of conversation list
            history_conv = session[i]
            count += len(history_conv["question"]) + len(history_conv["answer"])
            count_list.append(count)

        for c in count_list:
            if c > max_tokens:
                # pop first conversation
                session.pop(0)

    @staticmethod
    def clear_session(user_id):
        session_list[user_id] = []


def main():
    openai_bot = OpenAIBot()

    answer1 = openai_bot.reply(question="1+2的结果？", session_id=1)
    print(answer1)

    answer2 = openai_bot.reply(question="确定吗？不是等于5吗？", session_id=1)
    print(answer2)


if __name__ == '__main__':
    main()
